import sys
from PySide6.QtWidgets import QApplication, QDialog, QVBoxLayout, QTableView, QHBoxLayout, QCheckBox, QSpinBox, QLabel, \
    QPushButton, QFileDialog, QSplitter, QComboBox, QListWidget, QListWidgetItem, QMessageBox
from PySide6.QtCore import Qt
import pandas as pd
from ui import horizontal_visualization
import os
from core.pandas_model import PandasModel


class DataImportDialog(QDialog):
    def __init__(self, file_path, df, parent=None):
        super().__init__(parent)
        self.setWindowTitle("数据预览与导入")
        self.df = df
        self.file_path = file_path

        layout = QVBoxLayout(self)

        # 1. 预览区
        self.table = QTableView()
        self.model = PandasModel(df)
        self.table.setModel(self.model)
        layout.addWidget(QLabel("数据预览"))
        layout.addWidget(self.table)

        # 2. 字段选择区
        layout.addWidget(QLabel("选择要导入的字段（可多选）："))
        self.field_list = QListWidget()
        self.field_list.setSelectionMode(QListWidget.MultiSelection)
        for col in df.columns:
            item = QListWidgetItem(col)
            item.setSelected(True)
            self.field_list.addItem(item)
        layout.addWidget(self.field_list)

        # 3. 行范围选择区
        row_layout = QHBoxLayout()
        row_layout.addWidget(QLabel("导入行范围："))
        self.start_spin = QSpinBox()
        self.start_spin.setRange(0, len(df) - 1)
        self.start_spin.setValue(0)
        row_layout.addWidget(self.start_spin)
        row_layout.addWidget(QLabel("到"))
        self.end_spin = QSpinBox()
        self.end_spin.setRange(0, len(df) - 1)
        self.end_spin.setValue(len(df) - 1)
        row_layout.addWidget(self.end_spin)
        layout.addLayout(row_layout)

        # 4. 自动识别树型结构（如有"井号/井名/测井类型"等字段，自动分组）
        self.tree_hint = QLabel()
        self._update_tree_hint()
        layout.addWidget(self.tree_hint)

        # 5. 导入按钮
        btn = QPushButton("确认导入")
        btn.clicked.connect(self._on_accept)
        layout.addWidget(btn)

    def _on_accept(self):
        # 检查选择
        if not self.field_list.selectedItems():
            QMessageBox.warning(self, "提示", "请至少选择一个字段！")
            return
        if self.start_spin.value() > self.end_spin.value():
            QMessageBox.warning(self, "提示", "起始行不能大于结束行！")
            return
        self.accept()

    def get_processed_data(self):
        # 返回用户选择的字段和行范围的数据
        cols = [item.text() for item in self.field_list.selectedItems()]
        start, end = self.start_spin.value(), self.end_spin.value()
        return self.df.loc[start:end, cols]

    def _update_tree_hint(self):
        # 自动识别树型结构提示
        tree_fields = []
        for key in ["井号", "井名", "测井类型", "层号"]:
            if key in self.df.columns:
                tree_fields.append(key)
        if tree_fields:
            self.tree_hint.setText(f"建议树型分组字段：{' > '.join(tree_fields)}")
        else:
            self.tree_hint.setText("未检测到可自动分组的字段。")

    def _import_data(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择文件", "", "Excel Files (*.xlsx);;CSV Files (*.csv)")
        if not file_path:
            return
        df = pd.read_excel(file_path) if file_path.endswith('.xlsx') else pd.read_csv(file_path)
        dialog = DataImportDialog(file_path, df, parent=self)
        if dialog.exec_() == QDialog.Accepted:
            processed_df = dialog.get_processed_data()
            well_name = ...  # 用户输入或文件名提取
            self.data_model.add_data(well_name, processed_df)
            self._refresh_tree()

        splitter = QSplitter(Qt.Horizontal)
        self.data_preview = QTableView()
        splitter.addWidget(self.data_preview)
        self.visualization = horizontal_visualization(self)
        splitter.addWidget(self.visualization)
        self.setCentralWidget(splitter)

        self.x_combo = QComboBox()
        self.y_combo = QComboBox()
        # ... 用户选择后，调用 self._update_plot() 

    def import_all_txt_files(resources_dir, data_model, parent=None):
        for fname in os.listdir(resources_dir):
            if fname.endswith('.txt'):
                file_path = os.path.join(resources_dir, fname)
                # 自动检测分隔符和编码
                try:
                    df = pd.read_csv(file_path, sep=None, engine='python', encoding='utf-8')
                except Exception:
                    df = pd.read_csv(file_path, sep=None, engine='python', encoding='gbk')
                # 弹出预导入对话框
                dialog = DataImportDialog(file_path, df, parent=parent)
                if dialog.exec_() == QDialog.Accepted:
                    processed_df = dialog.get_processed_data()
                    well_name = os.path.splitext(fname)[0]# 提取井名
                    data_model.add_data(well_name, processed_df)

    def import_txt_file(file_path, data_model, parent=None):
        # 自动检测分隔符和编码
        try:
            df = pd.read_csv(file_path, sep=None, engine='python', encoding='utf-8')
        except Exception:
            df = pd.read_csv(file_path, sep=None, engine='python', encoding='gbk')
        dialog = DataImportDialog(file_path, df, parent=parent)
        if dialog.exec_() == QDialog.Accepted:
            processed_df = dialog.get_processed_data()
            well_name = os.path.splitext(os.path.basename(file_path))[0]
            # 检查数据库是否已有
            if not data_model.get_data(well_name):
                data_model.add_data(well_name, processed_df)
            # 刷新树、弹出表格窗口等
            parent._refresh_tree()
            parent.table_window.open_table(well_name) 